from emergency_plan.generate_stream import generate_plan_stream, generate_plan_stream_with_accumulation
from emergency_plan.query_rewrite import query_rewrite
from emergency_plan.recall import recall
from emergency_plan.rerank import rerank


class EmergencyPlanPipelineStream:
    """支持流式输出的应急预案生成pipeline"""
    
    def __init__(self, event_type: str, event_subtype: str, 
                 optimize_field: str = None, optimize_suggest: str = None, 
                 history: list = None):
        """
        初始化pipeline
        
        Args:
            event_type: 事件类型
            event_subtype: 事件子类
            optimize_field: 要优化的字段
            optimize_suggest: 优化建议
            history: 历史对话
        """
        self.event_type = event_type
        self.event_subtype = event_subtype
        self.optimize_field = optimize_field
        self.optimize_suggest = optimize_suggest
        self.history = history or []
    
    async def process_stream(self):
        """
        流式处理应急预案生成
        
        Yields:
            生成过程中的内容片段
        """
        # 构建查询
        query = f"事件类型：{self.event_type}，事件子类：{self.event_subtype}"
        
        print(f"🔍 EmergencyPlanPipelineStream开始处理")
        print(f"🔍 查询: {query}")
        print(f"🔍 优化字段: {self.optimize_field}")
        print(f"🔍 优化建议: {self.optimize_suggest}")
        print(f"🔍 历史对话数量: {len(self.history) if self.history else 0}")
        
        try:
            # 查询重写
            print(f"🔍 开始查询重写...")
            rewrite_result = await query_rewrite(query)
            print(f"✅ 查询重写完成: {rewrite_result}")
            
            # 召回相关文档
            print(f"🔍 开始文档召回...")
            recall_result = recall(rewrite_result)
            print(f"✅ 文档召回完成，召回数量: {len(recall_result)}")
            
            # 重排序
            print(f"🔍 开始文档重排序...")
            rerank_result = rerank(recall_result)
            print(f"✅ 文档重排序完成，重排序数量: {len(rerank_result)}")
            
            # 流式生成应急预案
            print(f"🔍 开始流式生成应急预案...")
            async for chunk in generate_plan_stream(
                rerank_result, 
                self.event_type, 
                self.event_subtype, 
                self.optimize_field, 
                self.optimize_suggest, 
                self.history
            ):
                yield chunk
            
            print(f"✅ 应急预案流式生成完成")
            
        except Exception as e:
            print(f"❌ EmergencyPlanPipelineStream处理失败: {e}")
            import traceback
            traceback.print_exc()
            yield f"处理失败: {e}"
    
    async def process_stream_with_accumulation(self):
        """
        流式处理应急预案生成并累积结果
        
        Yields:
            生成过程中的内容片段，最后包含解析结果
        """
        # 构建查询
        query = f"事件类型：{self.event_type}，事件子类：{self.event_subtype}"
        
        print(f"🔍 EmergencyPlanPipelineStream开始处理（带累积）")
        print(f"🔍 查询: {query}")
        print(f"🔍 优化字段: {self.optimize_field}")
        print(f"🔍 优化建议: {self.optimize_suggest}")
        print(f"🔍 历史对话数量: {len(self.history) if self.history else 0}")
        
        try:
            # 查询重写
            print(f"🔍 开始查询重写...")
            rewrite_result = await query_rewrite(query)
            print(f"✅ 查询重写完成: {rewrite_result}")
            
            # 召回相关文档
            print(f"🔍 开始文档召回...")
            recall_result = recall(rewrite_result)
            print(f"✅ 文档召回完成，召回数量: {len(recall_result)}")
            
            # 重排序
            print(f"🔍 开始文档重排序...")
            rerank_result = rerank(recall_result)
            print(f"✅ 文档重排序完成，重排序数量: {len(rerank_result)}")
            
            # 流式生成应急预案并累积结果
            print(f"🔍 开始流式生成应急预案（带累积）...")
            async for chunk in generate_plan_stream_with_accumulation(
                rerank_result, 
                self.event_type, 
                self.event_subtype, 
                self.optimize_field, 
                self.optimize_suggest, 
                self.history
            ):
                yield chunk
            
            print(f"✅ 应急预案流式生成完成（带累积）")
            
        except Exception as e:
            print(f"❌ EmergencyPlanPipelineStream处理失败: {e}")
            import traceback
            traceback.print_exc()
            yield f"处理失败: {e}" 